"""I3D head"""

import mindspore.nn as nn
from mindspore.ops import operations as P
from mindvideo.models.backbones.I3d import Unit3D
from mindvideo.common.utils.class_factory import ClassFactory, ModuleType

@ClassFactory.register(ModuleType.HEAD)
class I3dHead(nn.Cell):
    """
    I3dHead definition

    Args:
        num_classes (int): The number of classes.
        dropout_keep_prob (): A float value of prob.

    Returns:
        Tensor, output tensor.

    Examples:
        I3dHead(num_classes=400, dropout_keep_prob=0.5)
    """
    def __init__(self, num_classes=400, dropout_keep_prob=0.5):
        super(I3dHead, self).__init__()
        self._num_classes = num_classes
        self.dropout = nn.Dropout(dropout_keep_prob)
        self.logits = Unit3D(in_channels=384+384+128+128, out_channels=self._num_classes,
                             kernel_shape=(1, 1, 1),
                             padding=0,
                             activation_fn=None,
                             use_batch_norm=False,
                             has_bias=True)
        self.mean_op = P.ReduceMean()

    def construct(self, x):
        x = self.logits(self.dropout(x))
        x = x.squeeze(3).squeeze(3)
        x = self.mean_op(x, 2)
        return x
